{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# 1. Attention Mechanism\n",
    "## Attention\n",
    "### Scaled Dot-product Attention"
   ],
   "id": "b7400724fc107bdf"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:06.165287Z",
     "start_time": "2024-10-07T11:38:05.048662Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import os\n",
    "\n",
    "import torch\n",
    "from hf_transformers import BertModel, BertTokenizer\n",
    "\n",
    "model_name = 'bert-base-uncased'\n",
    "tokenizer = BertTokenizer.from_pretrained(model_name)\n",
    "model = BertModel.from_pretrained(model_name)\n"
   ],
   "id": "a0107e2a15846182",
   "outputs": [],
   "execution_count": 9
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:06.248259Z",
     "start_time": "2024-10-07T11:38:06.168886Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 编码文本\n",
    "inputs = tokenizer(\"Hello, my name is John\", return_tensors=\"pt\")\n",
    "outputs = model(**inputs)\n",
    "last_hidden_states = outputs.last_hidden_state  # 这是模型的输出"
   ],
   "id": "510e8c2fee704235",
   "outputs": [],
   "execution_count": 10
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:06.273020Z",
     "start_time": "2024-10-07T11:38:06.266700Z"
    }
   },
   "cell_type": "code",
   "source": "last_hidden_states",
   "id": "51536a47e80f1b65",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.0369,  0.6015, -0.3253,  ..., -0.2672,  0.2218,  0.7841],\n",
       "         [ 0.2551,  0.0815,  0.0220,  ...,  0.0708,  0.8906,  0.3853],\n",
       "         [-0.4311,  0.4639,  0.1747,  ..., -0.5010,  0.5965,  0.4691],\n",
       "         ...,\n",
       "         [-0.4714,  0.6831,  0.4085,  ..., -0.8256,  0.0825,  0.7012],\n",
       "         [-0.0656, -0.1387, -0.3846,  ...,  0.6084,  0.4762, -0.8739],\n",
       "         [ 0.7371,  0.0638, -0.4075,  ...,  0.0370, -0.6124, -0.1276]]],\n",
       "       grad_fn=<NativeLayerNormBackward0>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 11
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:06.737345Z",
     "start_time": "2024-10-07T11:38:06.369922Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torch import nn\n",
    "from hf_transformers import AutoConfig\n",
    "from hf_transformers import AutoTokenizer\n",
    "\n",
    "model_ckpt = \"bert-base-uncased\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_ckpt)"
   ],
   "id": "8181ceb2c9657b9f",
   "outputs": [],
   "execution_count": 12
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:06.756904Z",
     "start_time": "2024-10-07T11:38:06.753412Z"
    }
   },
   "cell_type": "code",
   "source": [
    "text = \"time flies like an arrow\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\", add_special_tokens=False)\n",
    "print(inputs.input_ids)"
   ],
   "id": "8379e26e66069580",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 2051, 10029,  2066,  2019,  8612]])\n"
     ]
    }
   ],
   "execution_count": 13
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:07.276999Z",
     "start_time": "2024-10-07T11:38:06.819979Z"
    }
   },
   "cell_type": "code",
   "source": [
    "config = AutoConfig.from_pretrained(model_ckpt)\n",
    "token_emb = nn.Embedding(config.vocab_size, config.hidden_size)\n",
    "print(token_emb)\n",
    "# input_ids 对应分词之后的 tokens 映射到的数字编号列表，而 attention_mask 则是用来标记哪些 tokens 是被填充的（这里“1”表示是原文，“0”表示是填充字符）。\n",
    "inputs_embeds = token_emb(inputs.input_ids)\n",
    "print(inputs_embeds.size())"
   ],
   "id": "6021f7fcbd9950df",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Embedding(30522, 768)\n",
      "torch.Size([1, 5, 768])\n"
     ]
    }
   ],
   "execution_count": 14
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:07.329942Z",
     "start_time": "2024-10-07T11:38:07.324561Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "from math import sqrt\n",
    "\n",
    "Q = K = V = inputs_embeds\n",
    "dim_k = K.size(-1)\n",
    "scores = torch.bmm(Q, K.transpose(1,2)) / sqrt(dim_k)\n",
    "print(scores.size())"
   ],
   "id": "115d8cfa669bbdd6",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 5])\n"
     ]
    }
   ],
   "execution_count": 15
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:07.432564Z",
     "start_time": "2024-10-07T11:38:07.374734Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "weights = F.softmax(scores, dim=-1)\n",
    "print(weights.sum(dim=-1))"
   ],
   "id": "514b28547ab98cf3",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 1., 1., 1., 1.]], grad_fn=<SumBackward1>)\n"
     ]
    }
   ],
   "execution_count": 16
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:07.514931Z",
     "start_time": "2024-10-07T11:38:07.511116Z"
    }
   },
   "cell_type": "code",
   "source": [
    "attn_outputs = torch.bmm(weights, V)\n",
    "print(attn_outputs.shape)"
   ],
   "id": "b525568080d04a83",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n"
     ]
    }
   ],
   "execution_count": 17
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:07.567344Z",
     "start_time": "2024-10-07T11:38:07.562849Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from math import sqrt\n",
    "\n",
    "def scaled_dot_product_attention(query, key, value, query_mask=None, key_mask=None, mask=None):\n",
    "    \"\"\"\n",
    "    计算 Scaled Dot-product Attention\n",
    "    :param query: \n",
    "    :param key: \n",
    "    :param value: \n",
    "    :param query_mask: 用于表示 query 中的有效位置，即哪些位置是需要计算注意力的。对于那些无效（如填充的部分），将它们标记为 0。\n",
    "    :param key_mask: 用于表示 key 中的有效位置，同样，0 表示无效或填充的部分。\n",
    "    :param mask: 用来覆盖掉无效的注意力分数（如填充的位置）。如果 mask 中某个位置为 0，那么对应的注意力分数将被设置为 -inf，这样在 softmax 操作后，注意力权重就会接近 0，从而忽略该位置。\n",
    "    :return: \n",
    "    \"\"\"\n",
    "    dim_k = query.size(-1)\n",
    "    scores = torch.bmm(query, key.transpose(1, 2)) / sqrt(dim_k)\n",
    "    if query_mask is not None and key_mask is not None:\n",
    "        # 通过矩阵相乘，mask 结果的形状为 [batch_size, query_length, key_length]，表示 query 中每个位置和 key 中每个位置之间的有效性对应关系。\n",
    "        mask = torch.bmm(query_mask.unsqueeze(-1), key_mask.unsqueeze(1))\n",
    "    if mask is not None:\n",
    "        scores = scores.masked_fill(mask == 0, -float(\"inf\"))\n",
    "    weights = F.softmax(scores, dim=-1)\n",
    "    return torch.bmm(weights, value)"
   ],
   "id": "f34450330995643d",
   "outputs": [],
   "execution_count": 18
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Multi-head Attention",
   "id": "b46e487dbe12d58b"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:07.615724Z",
     "start_time": "2024-10-07T11:38:07.609960Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torch import nn\n",
    "\n",
    "class AttentionHead(nn.Module):\n",
    "    def __init__(self, embed_dim, head_dim):\n",
    "        super().__init__()\n",
    "        # 每个头都会初始化三个独立的线性层，负责将 \n",
    "        # 序列映射到尺寸为 [batch_size, seq_len, head_dim] 的张量，其中 head_dim 是映射到的向量维度。\n",
    "        self.q = nn.Linear(embed_dim, head_dim)\n",
    "        self.k = nn.Linear(embed_dim, head_dim)\n",
    "        self.v = nn.Linear(embed_dim, head_dim)\n",
    "        # 实践中一般将 head_dim 设置为 embed_dim 的因数，这样 token 嵌入式表示的维度就可以保持不变，例如 BERT 有 12 个注意力头，因此每个头的维度被设置为 768 / 12 = 64\n",
    "\n",
    "    def forward(self, query, key, value, query_mask=None, key_mask=None, mask=None):\n",
    "        attn_outputs = scaled_dot_product_attention(\n",
    "            self.q(query), self.k(key), self.v(value), query_mask, key_mask, mask)\n",
    "        return attn_outputs"
   ],
   "id": "4463cf80c97e1ba3",
   "outputs": [],
   "execution_count": 19
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:07.641852Z",
     "start_time": "2024-10-07T11:38:07.636845Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        embed_dim = config.hidden_size\n",
    "        num_heads = config.num_attention_heads\n",
    "        head_dim = embed_dim // num_heads\n",
    "        self.heads = nn.ModuleList(\n",
    "            [AttentionHead(embed_dim, head_dim) for _ in range(num_heads)]\n",
    "        )\n",
    "        # 在拼接后还通过一个线性变换来生成最终的输出张量\n",
    "        self.output_linear = nn.Linear(embed_dim, embed_dim)\n",
    "\n",
    "    def forward(self, query, key, value, query_mask=None, key_mask=None, mask=None):\n",
    "        x = torch.cat([\n",
    "            h(query, key, value, query_mask, key_mask, mask) for h in self.heads\n",
    "        ], dim=-1)\n",
    "        x = self.output_linear(x)\n",
    "        return x"
   ],
   "id": "27d500c74e46257",
   "outputs": [],
   "execution_count": 20
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:08.481975Z",
     "start_time": "2024-10-07T11:38:07.659031Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from hf_transformers import AutoConfig\n",
    "from hf_transformers import AutoTokenizer\n",
    "\n",
    "model_ckpt = \"bert-base-uncased\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_ckpt)\n",
    "\n",
    "text = \"time flies like an arrow\"\n",
    "inputs = tokenizer(text, return_tensors=\"pt\", add_special_tokens=False)\n",
    "config = AutoConfig.from_pretrained(model_ckpt)\n",
    "token_emb = nn.Embedding(config.vocab_size, config.hidden_size)\n",
    "inputs_embeds = token_emb(inputs.input_ids)\n",
    "\n",
    "multihead_attn = MultiHeadAttention(config)\n",
    "query = key = value = inputs_embeds\n",
    "attn_output = multihead_attn(query, key, value)\n",
    "print(attn_output.size())"
   ],
   "id": "bd7e93c9ee27e04f",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n"
     ]
    }
   ],
   "execution_count": 21
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Transformer Encoder\n",
    "### The Feed-Forward Layer\n",
    "\n",
    "Transformer Encoder/Decoder 中的前馈子层实际上就是两层全连接神经网络，它单独地处理序列中的每一个词向量，也被称为 position-wise feed-forward layer。常见做法是让第一层的维度是词向量大小的 4 倍，然后以 GELU 作为激活函数。"
   ],
   "id": "5c738c48f8ca5270"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:08.667760Z",
     "start_time": "2024-10-07T11:38:08.586675Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class FeedForward(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.linear_1 = nn.Linear(config.hidden_size, config.intermediate_size)\n",
    "        self.linear_2 = nn.Linear(config.intermediate_size, config.hidden_size)\n",
    "        self.gelu = nn.GELU()\n",
    "        self.dropout = nn.Dropout(config.hidden_dropout_prob)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.linear_1(x)\n",
    "        x = self.gelu(x)\n",
    "        x = self.linear_2(x)\n",
    "        x = self.dropout(x)\n",
    "        return x\n",
    "    \n",
    "feed_forward = FeedForward(config)\n",
    "ff_outputs = feed_forward(attn_output)\n",
    "print(ff_outputs.size())"
   ],
   "id": "1ff963d971b66e83",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n"
     ]
    }
   ],
   "execution_count": 22
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### Layer Normalization\n",
    "Layer Normalization 负责将一批 (batch) 输入中的每一个都标准化为均值为零且具有单位方差；Skip Connections 则是将张量直接传递给模型的下一层而不进行处理，并将其添加到处理后的张量中。\n",
    "\n",
    "向 Transformer Encoder/Decoder 中添加 Layer Normalization 目前共有两种做法：\n",
    "\n",
    "![Layer Normalization](./images/LayerNormalization.png)\n",
    "\n",
    "- Post layer normalization：Transformer 论文中使用的方式，将 Layer normalization 放在 Skip Connections 之间。 但是因为梯度可能会发散，这种做法很难训练，还需要结合学习率预热 (learning rate warm-up) 等技巧；\n",
    "- Pre layer normalization：目前主流的做法，将 Layer Normalization 放置于 Skip Connections 的范围内。这种做法通常训练过程会更加稳定，并且不需要任何学习率预热。"
   ],
   "id": "1ddbdccc6cbaa613"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:08.804818Z",
     "start_time": "2024-10-07T11:38:08.765739Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class TransformerEncoderLayer(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.layer_norm_1 = nn.LayerNorm(config.hidden_size)\n",
    "        self.layer_norm_2 = nn.LayerNorm(config.hidden_size)\n",
    "        self.attention = MultiHeadAttention(config)\n",
    "        self.feed_forward = FeedForward(config)\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        # Apply layer normalization and then copy input into query, key, value\n",
    "        hidden_state = self.layer_norm_1(x)\n",
    "        # Apply attention with a skip connection\n",
    "        x = x + self.attention(hidden_state, hidden_state, hidden_state, mask=mask)\n",
    "        # Apply feed-forward layer with a skip connection\n",
    "        x = x + self.feed_forward(self.layer_norm_2(x))\n",
    "        return x\n",
    "    \n",
    "encoder_layer = TransformerEncoderLayer(config)\n",
    "print(inputs_embeds.shape)\n",
    "print(encoder_layer(inputs_embeds).size())"
   ],
   "id": "b165b393b8714a97",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n",
      "torch.Size([1, 5, 768])\n"
     ]
    }
   ],
   "execution_count": 23
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Positional Embeddings",
   "id": "8330e59c202013b6"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:09.026351Z",
     "start_time": "2024-10-07T11:38:08.915095Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Embeddings(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.token_embeddings = nn.Embedding(config.vocab_size,\n",
    "                                             config.hidden_size)\n",
    "        self.position_embeddings = nn.Embedding(config.max_position_embeddings,\n",
    "                                                config.hidden_size)\n",
    "        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=1e-12)\n",
    "        self.dropout = nn.Dropout()\n",
    "\n",
    "    def forward(self, input_ids):\n",
    "        # Create position IDs for input sequence\n",
    "        seq_length = input_ids.size(1)\n",
    "        position_ids = torch.arange(seq_length, dtype=torch.long).unsqueeze(0)\n",
    "        # Create token and position embeddings\n",
    "        token_embeddings = self.token_embeddings(input_ids)\n",
    "        position_embeddings = self.position_embeddings(position_ids)\n",
    "        # Combine token and position embeddings\n",
    "        embeddings = token_embeddings + position_embeddings\n",
    "        embeddings = self.layer_norm(embeddings)\n",
    "        embeddings = self.dropout(embeddings)\n",
    "        return embeddings\n",
    "\n",
    "embedding_layer = Embeddings(config)\n",
    "print(embedding_layer(inputs.input_ids).size())"
   ],
   "id": "5906dcac679a1d00",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n"
     ]
    }
   ],
   "execution_count": 24
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:09.645642Z",
     "start_time": "2024-10-07T11:38:09.073658Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class TransformerEncoder(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.embeddings = Embeddings(config)\n",
    "        self.layers = nn.ModuleList([TransformerEncoderLayer(config)\n",
    "                                     for _ in range(config.num_hidden_layers)])\n",
    "\n",
    "    def forward(self, x, mask=None):\n",
    "        x = self.embeddings(x)\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, mask=mask)\n",
    "        return x\n",
    "    \n",
    "encoder = TransformerEncoder(config)\n",
    "print(encoder(inputs.input_ids).size())"
   ],
   "id": "9fe36f4165b007b7",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 5, 768])\n"
     ]
    }
   ],
   "execution_count": 25
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Transformer Decoder\n",
    "Transformer Decoder 与 Encoder 最大的不同在于 Decoder 有两个注意力子层，如下图所示：\n",
    "\n",
    "![Transformer Decoder](./images/TransformerDecoder.png)\n",
    "\n",
    "- Masked multi-head self-attention layer：确保在每个时间步生成的词语仅基于过去的输出和当前预测的词，否则 Decoder 相当于作弊了；\n",
    "\n",
    "- Encoder-decoder attention layer：以解码器的中间表示作为 queries，对 encoder stack 的输出 key 和 value 向量执行 Multi-head Attention。通过这种方式，Encoder-Decoder Attention Layer 就可以学习到如何关联来自两个不同序列的词语，例如两种不同的语言。 解码器可以访问每个 block 中 Encoder 的 keys 和 values。"
   ],
   "id": "860504c45e0754b5"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:09.749729Z",
     "start_time": "2024-10-07T11:38:09.714924Z"
    }
   },
   "cell_type": "code",
   "source": [
    "seq_len = inputs.input_ids.size(-1)\n",
    "mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0)\n",
    "print(mask[0])"
   ],
   "id": "43aa2c5c38f051ca",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 0., 0., 0., 0.],\n",
      "        [1., 1., 0., 0., 0.],\n",
      "        [1., 1., 1., 0., 0.],\n",
      "        [1., 1., 1., 1., 0.],\n",
      "        [1., 1., 1., 1., 1.]])\n"
     ]
    }
   ],
   "execution_count": 26
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:38:09.850646Z",
     "start_time": "2024-10-07T11:38:09.826821Z"
    }
   },
   "cell_type": "code",
   "source": "scores.masked_fill(mask == 0, -float(\"inf\"))",
   "id": "5b9d81a720ea6df4",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[28.7794,    -inf,    -inf,    -inf,    -inf],\n",
       "         [-2.2343, 27.8819,    -inf,    -inf,    -inf],\n",
       "         [-0.9384,  0.9253, 29.4696,    -inf,    -inf],\n",
       "         [ 3.3696,  0.8230,  1.2350, 26.1508,    -inf],\n",
       "         [-0.9354,  0.3352,  0.5575, -0.5730, 26.4757]]],\n",
       "       grad_fn=<MaskedFillBackward0>)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 27
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# 2. Out of the box pipelines\n",
    "## Sentiment Analysis"
   ],
   "id": "b931d92766e64157"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:41:09.464084Z",
     "start_time": "2024-10-07T11:38:09.921619Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from hf_transformers import pipeline\n",
    "\n",
    "classifier = pipeline(\"sentiment-analysis\")\n",
    "result = classifier(\"I've been waiting for a HuggingFace course my whole life.\")\n",
    "print(result)\n",
    "results = classifier(\n",
    "  [\"I've been waiting for a HuggingFace course my whole life.\", \"I hate this so much!\"]\n",
    ")\n",
    "print(results)"
   ],
   "id": "57b062590cf8050f",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to distilbert/distilbert-base-uncased-finetuned-sst-2-english and revision 714eb0f (https://huggingface.co/distilbert/distilbert-base-uncased-finetuned-sst-2-english).\n",
      "Using a pipeline without specifying a model name and revision in production is not recommended.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "config.json:   0%|          | 0.00/629 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "dffd16090a424b32bacbd8e380fd8328"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/268M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "f7ddc977a80345609b6bd30994af1405"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/48.0 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "1061495a3db440e5a6511f26a15c16d4"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "vocab.txt:   0%|          | 0.00/232k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "02997be10ac14a8e8c671c3399887e80"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'label': 'POSITIVE', 'score': 0.9598049521446228}]\n",
      "[{'label': 'POSITIVE', 'score': 0.9598049521446228}, {'label': 'NEGATIVE', 'score': 0.9994558691978455}]\n"
     ]
    }
   ],
   "execution_count": 28
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Zero-Shot Classification\n",
    "零训练样本分类 pipeline 允许我们在不提供任何标注数据的情况下自定义分类标签。"
   ],
   "id": "93b224c122a3c293"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:53:55.770697Z",
     "start_time": "2024-10-07T11:41:09.497193Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from hf_transformers import pipeline\n",
    "\n",
    "classifier = pipeline(\"zero-shot-classification\")\n",
    "result = classifier(\n",
    "\"This is a course about the Transformers library\",\n",
    "candidate_labels=[\"education\", \"politics\", \"business\"],\n",
    ")\n",
    "print(result)"
   ],
   "id": "69d66b4204bb4812",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to facebook/bart-large-mnli and revision d7645e1 (https://huggingface.co/facebook/bart-large-mnli).\n",
      "Using a pipeline without specifying a model name and revision in production is not recommended.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "config.json:   0%|          | 0.00/1.15k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "4b9660522b304f8292f2eb46628b6c2b"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/1.63G [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "71a2f96a5d3e4a43bc960943b0ad679e"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/26.0 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "967505664f6f4ac08ee61ec497302b2a"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "vocab.json:   0%|          | 0.00/899k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "55e17fe5c3bc43339396f74157db1586"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "104341dd3cf648dfb1a6c025c5ecb654"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "dee86ad6aa2b41aa832d9f2b1a0b0821"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\transformers\\tokenization_utils_base.py:1617: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be deprecated in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'sequence': 'This is a course about the Transformers library', 'labels': ['education', 'business', 'politics'], 'scores': [0.8445996046066284, 0.11197374761104584, 0.04342668503522873]}\n"
     ]
    }
   ],
   "execution_count": 29
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Text Generation\n",
    "我们首先根据任务需要构建一个模板 (prompt)，然后将其送入到模型中来生成后续文本。注意，由于文本生成具有随机性，因此每次运行都会得到不同的结果。"
   ],
   "id": "32802a683566fa1b"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T11:57:39.824757Z",
     "start_time": "2024-10-07T11:54:02.081903Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from hf_transformers import pipeline\n",
    "\n",
    "generator = pipeline(\"text-generation\")\n",
    "results = generator(\"In this course, we will teach you how to\")\n",
    "print(results)\n",
    "results = generator(\n",
    "    \"In this course, we will teach you how to\",\n",
    "    num_return_sequences=2,\n",
    "    max_length=50\n",
    ") \n",
    "print(results)"
   ],
   "id": "e9e713281b1225aa",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to openai-community/gpt2 and revision 607a30d (https://huggingface.co/openai-community/gpt2).\n",
      "Using a pipeline without specifying a model name and revision in production is not recommended.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "config.json:   0%|          | 0.00/665 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "3d86cbb1505d4f28a31e52cbd6cefb65"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/548M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "fbd205703f46412f9f37834479e4ed62"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "generation_config.json:   0%|          | 0.00/124 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "a3de0e2c1ca547cb829eab6f9e1f9bfa"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/26.0 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "bb5352b9c6a84b05a0fb929459da80b5"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "vocab.json:   0%|          | 0.00/1.04M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b30aa48efce443d78709d4306f886a5f"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "aed4a23ad8364b1babc49954d1ccfc13"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "472cbbf5ebc84d39a6595e0db51c37c3"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n",
      "Truncation was not explicitly activated but `max_length` is provided a specific value, please use `truncation=True` to explicitly truncate examples to max length. Defaulting to 'longest_first' truncation strategy. If you encode pairs of sequences (GLUE-style) with the tokenizer you can select this strategy more precisely by providing a specific strategy to `truncation`.\n",
      "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'generated_text': 'In this course, we will teach you how to create a secure connection with Android.\\n\\nLearn how to make an encrypted connection with OSS, HTTP, and HTTPS.'}]\n",
      "[{'generated_text': 'In this course, we will teach you how to create and implement user-defined properties, provide APIs for creating, building, modifying and interacting with the assets in your project from top to bottom. This class assumes you have a basic knowledge of how to'}, {'generated_text': 'In this course, we will teach you how to create high quality sound effects using real-time dynamics (Rendering) and stereo rendering. Using the following techniques, we will be able to develop a soundstage through using different sound processing techniques by'}]\n"
     ]
    }
   ],
   "execution_count": 30
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T12:01:01.481310Z",
     "start_time": "2024-10-07T11:57:39.854185Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from hf_transformers import pipeline\n",
    "\n",
    "generator = pipeline(\"text-generation\", model=\"distilgpt2\")\n",
    "results = generator(\n",
    "    \"In this course, we will teach you how to\",\n",
    "    max_length=30,\n",
    "    num_return_sequences=2,\n",
    ")\n",
    "print(results)"
   ],
   "id": "20f912f9d0df86ac",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "config.json:   0%|          | 0.00/762 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e1f8d907780e4c8bae9efab7674a5a51"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/353M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "10962725d9234914add73a15a413aa44"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "generation_config.json:   0%|          | 0.00/124 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "7d7f0667e33b405cb0caca1e9f6a8cd3"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/26.0 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "45b079a9e49547a9a227b87ee46642c3"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "vocab.json:   0%|          | 0.00/1.04M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "cfcb146d914e4c789e14fdb27d71b71d"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8985dae803de4ab3849df910fa57331d"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "cb8b5d4846de48fe99eba062817c75bb"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Truncation was not explicitly activated but `max_length` is provided a specific value, please use `truncation=True` to explicitly truncate examples to max length. Defaulting to 'longest_first' truncation strategy. If you encode pairs of sequences (GLUE-style) with the tokenizer you can select this strategy more precisely by providing a specific strategy to `truncation`.\n",
      "Setting `pad_token_id` to `eos_token_id`:None for open-end generation.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'generated_text': 'In this course, we will teach you how to convert, convert, and help you discover how to convert into your life through a new format for you'}, {'generated_text': 'In this course, we will teach you how to use one of the most powerful in-depth and effective non-fiction websites.\\n\\n\\n\\n'}]\n"
     ]
    }
   ],
   "execution_count": 31
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T12:04:34.928424Z",
     "start_time": "2024-10-07T12:01:01.513037Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from hf_transformers import pipeline\n",
    "\n",
    "generator = pipeline(\"text-generation\", model=\"uer/gpt2-chinese-poem\")\n",
    "results = generator(\n",
    "    \"[CLS] 万 叠 春 山 积 雨 晴 ，\",\n",
    "    max_length=40,\n",
    "    num_return_sequences=2,\n",
    ")\n",
    "print(results)"
   ],
   "id": "ec389f76be70a6ce",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "config.json:   0%|          | 0.00/577 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ecfc9ccbfae8464ca03472e811e62150"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/425M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e58ce3805aca407da24f4265c0b0f32e"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/271 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "b01e1f832e294d7c8a35f66b874eee44"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "vocab.txt:   0%|          | 0.00/115k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "a9ced1864800430f8ce747e3d857f4e1"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/112 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "d373502c6738441ebb1c0355e8f03fcd"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Truncation was not explicitly activated but `max_length` is provided a specific value, please use `truncation=True` to explicitly truncate examples to max length. Defaulting to 'longest_first' truncation strategy. If you encode pairs of sequences (GLUE-style) with the tokenizer you can select this strategy more precisely by providing a specific strategy to `truncation`.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'generated_text': '[CLS] 万 叠 春 山 积 雨 晴 ， 晓 日 照 松 。 闻 寒 鹊 噪 声 急 ， 风 折 珊 瑚 枝 ， 谷 应 涛 白 。 然 花 落 ，'}, {'generated_text': '[CLS] 万 叠 春 山 积 雨 晴 ， 媚 不 可 收 。 何 似 东 君 景 与 西 ， 气 势 磅 礴 蟠 天 浮 。 高 柳 阴 连 古 渡 头 ，'}]\n"
     ]
    }
   ],
   "execution_count": 32
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Fill Mask\n",
    "给定一段部分词语被遮盖掉 (masked) 的文本，使用预训练模型来预测能够填充这些位置的词语。"
   ],
   "id": "49d577491c1f3e65"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T12:07:59.644451Z",
     "start_time": "2024-10-07T12:04:34.962789Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from hf_transformers import pipeline\n",
    "\n",
    "unmasker = pipeline(\"fill-mask\")\n",
    "results = unmasker(\"This course will teach you all about <mask> models.\", top_k=2)\n",
    "print(results)"
   ],
   "id": "46b96680f0a7c8ba",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to distilbert/distilroberta-base and revision fb53ab8 (https://huggingface.co/distilbert/distilroberta-base).\n",
      "Using a pipeline without specifying a model name and revision in production is not recommended.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "config.json:   0%|          | 0.00/480 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "10d4b4ee74cd41ee98406c3afae6b410"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/331M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "1c08f2a6eb674c36b1d5cdc6b7fa0c2c"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilbert/distilroberta-base were not used when initializing RobertaForMaskedLM: ['roberta.pooler.dense.bias', 'roberta.pooler.dense.weight']\n",
      "- This IS expected if you are initializing RobertaForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing RobertaForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/25.0 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "07ccae3df6a8477380b660bbdcaf25e3"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "vocab.json:   0%|          | 0.00/899k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c0587709b0d44974b8743cfdcbc4ba85"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ce1f8882e9104d5b8dfac1c58ba07bbc"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "924ff427b1a8429586a21688103a147d"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'score': 0.19198425114154816, 'token': 30412, 'token_str': ' mathematical', 'sequence': 'This course will teach you all about mathematical models.'}, {'score': 0.042092215269804, 'token': 38163, 'token_str': ' computational', 'sequence': 'This course will teach you all about computational models.'}]\n"
     ]
    }
   ],
   "execution_count": 33
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Named Entity Recognition\n",
    "命名实体识别 (NER) pipeline 负责从文本中抽取出指定类型的实体，例如人物、地点、组织等等。"
   ],
   "id": "a0281e3adace3acd"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T12:25:28.850240Z",
     "start_time": "2024-10-07T12:07:59.681305Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from hf_transformers import pipeline\n",
    "\n",
    "ner = pipeline(\"ner\", grouped_entities=True)\n",
    "results = ner(\"My name is Sylvain and I work at Hugging Face in Brooklyn.\")\n",
    "print(results)"
   ],
   "id": "7afbbf52d9471168",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to dbmdz/bert-large-cased-finetuned-conll03-english and revision 4c53496 (https://huggingface.co/dbmdz/bert-large-cased-finetuned-conll03-english).\n",
      "Using a pipeline without specifying a model name and revision in production is not recommended.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "config.json:   0%|          | 0.00/998 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "3a9ecf62ace9477894cc3c9b705be423"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/1.33G [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "d6f8e3c0829144798eea5df7552ec044"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at dbmdz/bert-large-cased-finetuned-conll03-english were not used when initializing BertForTokenClassification: ['bert.pooler.dense.bias', 'bert.pooler.dense.weight']\n",
      "- This IS expected if you are initializing BertForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/60.0 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "609cf9d4f335431dad47bbaa1fce89cf"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "vocab.txt:   0%|          | 0.00/213k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9951b48dcf0d43c99b5070720ad1ad42"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\transformers\\pipelines\\token_classification.py:170: UserWarning: `grouped_entities` is deprecated and will be removed in version v5.0.0, defaulted to `aggregation_strategy=\"AggregationStrategy.SIMPLE\"` instead.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'entity_group': 'PER', 'score': 0.9981694, 'word': 'Sylvain', 'start': 11, 'end': 18}, {'entity_group': 'ORG', 'score': 0.9796019, 'word': 'Hugging Face', 'start': 33, 'end': 45}, {'entity_group': 'LOC', 'score': 0.9932106, 'word': 'Brooklyn', 'start': 49, 'end': 57}]\n"
     ]
    }
   ],
   "execution_count": 34
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Question Answering\n",
    "自动问答 pipeline 可以根据给定的上下文回答问题"
   ],
   "id": "30e5b947c4c0ac26"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T12:26:49.238930Z",
     "start_time": "2024-10-07T12:25:29.042038Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from hf_transformers import pipeline\n",
    "\n",
    "question_answerer = pipeline(\"question-answering\")\n",
    "answer = question_answerer(\n",
    "    question=\"Where do I work?\",\n",
    "    context=\"My name is Sylvain and I work at Hugging Face in Brooklyn\",\n",
    ")\n",
    "print(answer)"
   ],
   "id": "8444ccd99d13bb62",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to distilbert/distilbert-base-cased-distilled-squad and revision 564e9b5 (https://huggingface.co/distilbert/distilbert-base-cased-distilled-squad).\n",
      "Using a pipeline without specifying a model name and revision in production is not recommended.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "config.json:   0%|          | 0.00/473 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "053cbb626f4149dbb14b289a25a748fe"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/261M [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "1a5f213d49a14b9abdcdd066b0d84c9c"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/49.0 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "6577cf19c4ea43dfa571c16d78025ef5"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "vocab.txt:   0%|          | 0.00/213k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ee0c579213984ffbb2d472051eed890b"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/436k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e34c1cff8c5a47c384c1d1852ab00b54"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'score': 0.6949766278266907, 'start': 33, 'end': 45, 'answer': 'Hugging Face'}\n"
     ]
    }
   ],
   "execution_count": 35
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Summarization\n",
    "自动摘要 pipeline 旨在将长文本压缩成短文本，并且还要尽可能保留原文的主要信息，例如："
   ],
   "id": "efe5d89db646bc93"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T12:38:22.743107Z",
     "start_time": "2024-10-07T12:26:49.299223Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from hf_transformers import pipeline\n",
    "\n",
    "summarizer = pipeline(\"summarization\")\n",
    "results = summarizer(\n",
    "    \"\"\"\n",
    "    America has changed dramatically during recent years. Not only has the number of \n",
    "    graduates in traditional engineering disciplines such as mechanical, civil, \n",
    "    electrical, chemical, and aeronautical engineering declined, but in most of \n",
    "    the premier American universities engineering curricula now concentrate on \n",
    "    and encourage largely the study of engineering science. As a result, there \n",
    "    are declining offerings in engineering subjects dealing with infrastructure, \n",
    "    the environment, and related issues, and greater concentration on high \n",
    "    technology subjects, largely supporting increasingly complex scientific \n",
    "    developments. While the latter is important, it should not be at the expense \n",
    "    of more traditional engineering.\n",
    "\n",
    "    Rapidly developing economies such as China and India, as well as other \n",
    "    industrial countries in Europe and Asia, continue to encourage and advance \n",
    "    the teaching of engineering. Both China and India, respectively, graduate \n",
    "    six and eight times as many traditional engineers as does the United States. \n",
    "    Other industrial countries at minimum maintain their output, while America \n",
    "    suffers an increasingly serious decline in the number of engineering graduates \n",
    "    and a lack of well-educated engineers.\n",
    "    \"\"\"\n",
    ")\n",
    "print(results)"
   ],
   "id": "a4da51ad7891a653",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to sshleifer/distilbart-cnn-12-6 and revision a4f8f3e (https://huggingface.co/sshleifer/distilbart-cnn-12-6).\n",
      "Using a pipeline without specifying a model name and revision in production is not recommended.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "config.json:   0%|          | 0.00/1.80k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "9e55a8137dfa4362a48ba22b4d65489d"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "pytorch_model.bin:   0%|          | 0.00/1.22G [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "edb21d58b3f141888735cf36166c9e85"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/26.0 [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "16885655e0764292936beed8f0f87fa0"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "vocab.json:   0%|          | 0.00/899k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "bfe4fbb235354a1f9b79a7f00e2b8c4b"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "5e353e731ec34635ab56224688501f3a"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'summary_text': ' America has changed dramatically during recent years . The number of engineering graduates in the U.S. has declined in traditional engineering disciplines such as mechanical, civil, electrical, chemical, and aeronautical engineering . Rapidly developing economies such as China and India, as well as other industrial countries in Europe and Asia, continue to encourage and advance engineering .'}]\n"
     ]
    }
   ],
   "execution_count": 36
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## What's going on behind these pipelines?\n",
    "实际上它的背后经过了三个步骤：\n",
    "\n",
    "1. 预处理 (preprocessing)，将原始文本转换为模型可以接受的输入格式；\n",
    "2. 将处理好的输入送入模型；\n",
    "3. 对模型的输出进行后处理 (postprocessing)，将其转换为人类方便阅读的格式。\n",
    "\n",
    "![pipelines](./images/pipelines.png)"
   ],
   "id": "5bc9db98aaea9728"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "# 3. Model and Tokenizer\n",
    "## Padding\n",
    "按批输入多段文本产生的一个直接问题就是：batch 中的文本有长有短，而输入张量必须是严格的二维矩形，维度为 (batch size, sequence size)，即每一段文本编码后的 token IDs 数量必须一样多。\n",
    "## Attention Mask\n",
    "Attention Mask 是一个尺寸与 input IDs 完全相同，且仅由 0 和 1 组成的张量，0 表示对应位置的 token 是填充符，不参与计算。当然，一些特殊的模型结构也会借助 Attention Mask 来遮蔽掉指定的 tokens。\n",
    "\n",
    "# 4. Necessary Pytorch Knowledge\n",
    "## Pytorch Basics\n",
    "### Tensor\n",
    "直接在 GPU 中创建张量或者将张量送入到 GPU 中\n"
   ],
   "id": "239d3f551b658495"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T12:38:23.101794Z",
     "start_time": "2024-10-07T12:38:22.846475Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "print(torch.rand(2, 3).cuda())\n",
    "print(torch.rand(2, 3, device=\"cuda\"))\n",
    "print(torch.rand(2, 3).to(\"cuda\"))"
   ],
   "id": "e582d2347515817f",
   "outputs": [
    {
     "ename": "AssertionError",
     "evalue": "Torch not compiled with CUDA enabled",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mAssertionError\u001B[0m                            Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[37], line 2\u001B[0m\n\u001B[0;32m      1\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtorch\u001B[39;00m\n\u001B[1;32m----> 2\u001B[0m \u001B[38;5;28mprint\u001B[39m(torch\u001B[38;5;241m.\u001B[39mrand(\u001B[38;5;241m2\u001B[39m, \u001B[38;5;241m3\u001B[39m)\u001B[38;5;241m.\u001B[39mcuda())\n\u001B[0;32m      3\u001B[0m \u001B[38;5;28mprint\u001B[39m(torch\u001B[38;5;241m.\u001B[39mrand(\u001B[38;5;241m2\u001B[39m, \u001B[38;5;241m3\u001B[39m, device\u001B[38;5;241m=\u001B[39m\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mcuda\u001B[39m\u001B[38;5;124m\"\u001B[39m))\n\u001B[0;32m      4\u001B[0m \u001B[38;5;28mprint\u001B[39m(torch\u001B[38;5;241m.\u001B[39mrand(\u001B[38;5;241m2\u001B[39m, \u001B[38;5;241m3\u001B[39m)\u001B[38;5;241m.\u001B[39mto(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mcuda\u001B[39m\u001B[38;5;124m\"\u001B[39m))\n",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\torch\\cuda\\__init__.py:284\u001B[0m, in \u001B[0;36m_lazy_init\u001B[1;34m()\u001B[0m\n\u001B[0;32m    279\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mRuntimeError\u001B[39;00m(\n\u001B[0;32m    280\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mCannot re-initialize CUDA in forked subprocess. To use CUDA with \u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[0;32m    281\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mmultiprocessing, you must use the \u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mspawn\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124m start method\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[0;32m    282\u001B[0m     )\n\u001B[0;32m    283\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28mhasattr\u001B[39m(torch\u001B[38;5;241m.\u001B[39m_C, \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m_cuda_getDeviceCount\u001B[39m\u001B[38;5;124m\"\u001B[39m):\n\u001B[1;32m--> 284\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mAssertionError\u001B[39;00m(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mTorch not compiled with CUDA enabled\u001B[39m\u001B[38;5;124m\"\u001B[39m)\n\u001B[0;32m    285\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m _cudart \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[0;32m    286\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mAssertionError\u001B[39;00m(\n\u001B[0;32m    287\u001B[0m         \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mlibcudart functions unavailable. It looks like you have a broken build?\u001B[39m\u001B[38;5;124m\"\u001B[39m\n\u001B[0;32m    288\u001B[0m     )\n",
      "\u001B[1;31mAssertionError\u001B[0m: Torch not compiled with CUDA enabled"
     ]
    }
   ],
   "execution_count": 37
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Automatic Differentiation",
   "id": "3bed42ffbfc767ac"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T14:05:13.852232Z",
     "start_time": "2024-10-07T14:05:13.793392Z"
    }
   },
   "cell_type": "code",
   "source": [
    "x = torch.tensor([2.], requires_grad=True)\n",
    "y = torch.tensor([3.], requires_grad=True)\n",
    "z = (x + y) * (y - 2)\n",
    "print(z)\n",
    "z.backward()\n",
    "print(x.grad, y.grad)"
   ],
   "id": "2c81f94250cd3081",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([5.], grad_fn=<MulBackward0>)\n",
      "tensor([1.]) tensor([6.])\n"
     ]
    }
   ],
   "execution_count": 38
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### Adjust the Tensor Shape\n",
    "\n",
    "- 形状转换 view 将张量转换为新的形状，需要保证总的元素个数不变，例如：\n",
    "```python\n",
    ">>> x = torch.tensor([1, 2, 3, 4, 5, 6])\n",
    ">>> print(x, x.shape)\n",
    "tensor([1, 2, 3, 4, 5, 6]) torch.Size([6])\n",
    ">>> x.view(2, 3) # shape adjusted to (2, 3)\n",
    "tensor([[1, 2, 3],\n",
    "        [4, 5, 6]])\n",
    ">>> x.view(3, 2) # shape adjusted to (3, 2)\n",
    "tensor([[1, 2],\n",
    "        [3, 4],\n",
    "        [5, 6]])\n",
    ">>> x.view(-1, 3) # -1 means automatic inference\n",
    "tensor([[1, 2, 3],\n",
    "        [4, 5, 6]])\n",
    "```\n",
    "进行 view 操作的张量必须是连续的 (contiguous)，可以调用 is_conuous 来判断张量是否连续；如果非连续，需要先通过 contiguous 函数将其变为连续的。也可以直接调用 Pytorch 新提供的 reshape 函数，它与 view 功能几乎一致，并且能够自动处理非连续张量。\n",
    "- 转置 transpose 交换张量中的两个维度，参数为相应的维度：\n",
    "```python\n",
    ">>> x = torch.tensor([[1, 2, 3], [4, 5, 6]])\n",
    ">>> x\n",
    "tensor([[1, 2, 3],\n",
    "        [4, 5, 6]])\n",
    ">>> x.transpose(0, 1)\n",
    "tensor([[1, 4],\n",
    "        [2, 5],\n",
    "        [3, 6]])\n",
    "```\n",
    "- 交换维度 permute 与 transpose 函数每次只能交换两个维度不同，permute 可以直接设置新的维度排列方式：\n",
    "```python\n",
    ">>> x = torch.tensor([[[1, 2, 3], [4, 5, 6]]])\n",
    ">>> print(x, x.shape)\n",
    "tensor([[[1, 2, 3],\n",
    "         [4, 5, 6]]]) torch.Size([1, 2, 3])\n",
    ">>> x = x.permute(2, 0, 1)\n",
    ">>> print(x, x.shape)\n",
    "tensor([[[1, 4]],\n",
    "  \n",
    "        [[2, 5]],\n",
    "  \n",
    "        [[3, 6]]]) torch.Size([3, 1, 2])\n",
    "```\n",
    "\n",
    "### Dimensionality Reduction and Dimensionality Increase\n",
    "有时为了计算需要对一个张量进行降维或升维。例如神经网络通常只接受一个批次 (batch) 的样例作为输入，如果只有 1 个输入样例，就需要手工添加一个 batch 维度。具体地：\n",
    "\n",
    "- 升维 torch.unsqueeze(input, dim, out=None) 在输入张量的 dim 位置插入一维，与索引一样，dim 值也可以为负数；\n",
    "- 降维 torch.squeeze(input, dim=None, out=None) 在不指定 dim 时，张量中所有形状为 1 的维度都会被删除，例如 (A,1,B,1,C) 会变成 (A,B,C)；当给定 dim 时，只会删除给定的维度（形状必须为 1），例如对于 (A,1,B)，squeeze(input, dim=0) 会保持张量不变，只有 squeeze(input, dim=1) 形状才会变成 (A,B)。\n",
    "下面是一些示例：\n",
    "```python\n",
    ">>> a = torch.tensor([1, 2, 3, 4])\n",
    ">>> a.shape\n",
    "torch.Size([4])\n",
    ">>> b = torch.unsqueeze(a, dim=0)\n",
    ">>> print(b, b.shape)\n",
    "tensor([[1, 2, 3, 4]]) torch.Size([1, 4])\n",
    ">>> b = a.unsqueeze(dim=0) # another way to unsqueeze tensor\n",
    ">>> print(b, b.shape)\n",
    "tensor([[1, 2, 3, 4]]) torch.Size([1, 4])\n",
    ">>> c = b.squeeze()\n",
    ">>> print(c, c.shape)\n",
    "tensor([1, 2, 3, 4]) torch.Size([4])\n",
    "```\n"
   ],
   "id": "fe593346ca9df7c4"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Loading Data\n",
    "### Dataset\n",
    "数据集负责存储数据样本，所有的数据集类都必须继承自 Dataset 或 IterableDataset。具体地，Pytorch 支持两种形式的数据集：\n",
    "\n",
    "- 映射型 (Map-style) 数据集\n",
    "\n",
    "继承自 Dataset 类，表示一个从索引到样本的映射（索引可以不是整数），这样我们就可以方便地通过 dataset[idx] 来访问指定索引的样本。这也是目前最常见的数据集类型。映射型数据集必须实现 __getitem__() 函数，其负责根据指定的 key 返回对应的样本。一般还会实现 __len__() 用于返回数据集的大小。\n",
    "\n",
    "> DataLoader 在默认情况下会创建一个生成整数索引的索引采样器 (sampler) 用于遍历数据集。因此，如果我们加载的是一个非整数索引的映射型数据集，还需要手工定义采样器。\n",
    "\n",
    "- 迭代型 (Iterable-style) 数据集\n",
    "\n",
    "继承自 IterableDataset，表示可迭代的数据集，它可以通过 iter(dataset) 以数据流 (steam) 的形式访问，适用于访问超大数据集或者远程服务器产生的数据。 迭代型数据集必须实现 __iter__() 函数，用于返回一个样本迭代器 (iterator)。\n",
    "\n",
    "> 注意：如果在 DataLoader 中开启多进程（num_workers > 0），那么在加载迭代型数据集时必须进行专门的设置，否则会重复访问样本。例如："
   ],
   "id": "20999cf93f021528"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T14:05:19.127118Z",
     "start_time": "2024-10-07T14:05:13.898030Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torch.utils.data import IterableDataset, DataLoader\n",
    "\n",
    "class MyIterableDataset(IterableDataset):\n",
    "\n",
    "    def __init__(self, start, end):\n",
    "        super(MyIterableDataset).__init__()\n",
    "        assert end > start\n",
    "        self.start = start\n",
    "        self.end = end\n",
    "\n",
    "    def __iter__(self):\n",
    "        return iter(range(self.start, self.end))\n",
    "\n",
    "ds = MyIterableDataset(start=3, end=7) # [3, 4, 5, 6]\n",
    "# Single-process loading\n",
    "print(list(DataLoader(ds, num_workers=0)))\n",
    "# Directly doing multi-process loading\n",
    "print(list(DataLoader(ds, num_workers=2)))"
   ],
   "id": "570275c689577c85",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[tensor([3]), tensor([4]), tensor([5]), tensor([6])]\n"
     ]
    },
    {
     "ename": "RuntimeError",
     "evalue": "DataLoader worker (pid(s) 15976, 6368) exited unexpectedly",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mEmpty\u001B[0m                                     Traceback (most recent call last)",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:1133\u001B[0m, in \u001B[0;36m_MultiProcessingDataLoaderIter._try_get_data\u001B[1;34m(self, timeout)\u001B[0m\n\u001B[0;32m   1132\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m-> 1133\u001B[0m     data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_data_queue\u001B[38;5;241m.\u001B[39mget(timeout\u001B[38;5;241m=\u001B[39mtimeout)\n\u001B[0;32m   1134\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m (\u001B[38;5;28;01mTrue\u001B[39;00m, data)\n",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\multiprocessing\\queues.py:114\u001B[0m, in \u001B[0;36mQueue.get\u001B[1;34m(self, block, timeout)\u001B[0m\n\u001B[0;32m    113\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_poll(timeout):\n\u001B[1;32m--> 114\u001B[0m         \u001B[38;5;28;01mraise\u001B[39;00m Empty\n\u001B[0;32m    115\u001B[0m \u001B[38;5;28;01melif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_poll():\n",
      "\u001B[1;31mEmpty\u001B[0m: ",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001B[1;31mRuntimeError\u001B[0m                              Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[39], line 18\u001B[0m\n\u001B[0;32m     16\u001B[0m \u001B[38;5;28mprint\u001B[39m(\u001B[38;5;28mlist\u001B[39m(DataLoader(ds, num_workers\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m0\u001B[39m)))\n\u001B[0;32m     17\u001B[0m \u001B[38;5;66;03m# Directly doing multi-process loading\u001B[39;00m\n\u001B[1;32m---> 18\u001B[0m \u001B[38;5;28mprint\u001B[39m(\u001B[38;5;28mlist\u001B[39m(DataLoader(ds, num_workers\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m2\u001B[39m)))\n",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:631\u001B[0m, in \u001B[0;36m_BaseDataLoaderIter.__next__\u001B[1;34m(self)\u001B[0m\n\u001B[0;32m    628\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_sampler_iter \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[0;32m    629\u001B[0m     \u001B[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001B[39;00m\n\u001B[0;32m    630\u001B[0m     \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_reset()  \u001B[38;5;66;03m# type: ignore[call-arg]\u001B[39;00m\n\u001B[1;32m--> 631\u001B[0m data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_next_data()\n\u001B[0;32m    632\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_num_yielded \u001B[38;5;241m+\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;241m1\u001B[39m\n\u001B[0;32m    633\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_dataset_kind \u001B[38;5;241m==\u001B[39m _DatasetKind\u001B[38;5;241m.\u001B[39mIterable \u001B[38;5;129;01mand\u001B[39;00m \\\n\u001B[0;32m    634\u001B[0m         \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_IterableDataset_len_called \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m \u001B[38;5;129;01mand\u001B[39;00m \\\n\u001B[0;32m    635\u001B[0m         \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_num_yielded \u001B[38;5;241m>\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_IterableDataset_len_called:\n",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:1329\u001B[0m, in \u001B[0;36m_MultiProcessingDataLoaderIter._next_data\u001B[1;34m(self)\u001B[0m\n\u001B[0;32m   1326\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_process_data(data)\n\u001B[0;32m   1328\u001B[0m \u001B[38;5;28;01massert\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_shutdown \u001B[38;5;129;01mand\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_tasks_outstanding \u001B[38;5;241m>\u001B[39m \u001B[38;5;241m0\u001B[39m\n\u001B[1;32m-> 1329\u001B[0m idx, data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_get_data()\n\u001B[0;32m   1330\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_tasks_outstanding \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;241m1\u001B[39m\n\u001B[0;32m   1331\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_dataset_kind \u001B[38;5;241m==\u001B[39m _DatasetKind\u001B[38;5;241m.\u001B[39mIterable:\n\u001B[0;32m   1332\u001B[0m     \u001B[38;5;66;03m# Check for _IterableDatasetStopIteration\u001B[39;00m\n",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:1295\u001B[0m, in \u001B[0;36m_MultiProcessingDataLoaderIter._get_data\u001B[1;34m(self)\u001B[0m\n\u001B[0;32m   1291\u001B[0m     \u001B[38;5;66;03m# In this case, `self._data_queue` is a `queue.Queue`,. But we don't\u001B[39;00m\n\u001B[0;32m   1292\u001B[0m     \u001B[38;5;66;03m# need to call `.task_done()` because we don't use `.join()`.\u001B[39;00m\n\u001B[0;32m   1293\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m   1294\u001B[0m     \u001B[38;5;28;01mwhile\u001B[39;00m \u001B[38;5;28;01mTrue\u001B[39;00m:\n\u001B[1;32m-> 1295\u001B[0m         success, data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_try_get_data()\n\u001B[0;32m   1296\u001B[0m         \u001B[38;5;28;01mif\u001B[39;00m success:\n\u001B[0;32m   1297\u001B[0m             \u001B[38;5;28;01mreturn\u001B[39;00m data\n",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:1146\u001B[0m, in \u001B[0;36m_MultiProcessingDataLoaderIter._try_get_data\u001B[1;34m(self, timeout)\u001B[0m\n\u001B[0;32m   1144\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mlen\u001B[39m(failed_workers) \u001B[38;5;241m>\u001B[39m \u001B[38;5;241m0\u001B[39m:\n\u001B[0;32m   1145\u001B[0m     pids_str \u001B[38;5;241m=\u001B[39m \u001B[38;5;124m'\u001B[39m\u001B[38;5;124m, \u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;241m.\u001B[39mjoin(\u001B[38;5;28mstr\u001B[39m(w\u001B[38;5;241m.\u001B[39mpid) \u001B[38;5;28;01mfor\u001B[39;00m w \u001B[38;5;129;01min\u001B[39;00m failed_workers)\n\u001B[1;32m-> 1146\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mRuntimeError\u001B[39;00m(\u001B[38;5;124mf\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mDataLoader worker (pid(s) \u001B[39m\u001B[38;5;132;01m{\u001B[39;00mpids_str\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m) exited unexpectedly\u001B[39m\u001B[38;5;124m'\u001B[39m) \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01me\u001B[39;00m\n\u001B[0;32m   1147\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(e, queue\u001B[38;5;241m.\u001B[39mEmpty):\n\u001B[0;32m   1148\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m (\u001B[38;5;28;01mFalse\u001B[39;00m, \u001B[38;5;28;01mNone\u001B[39;00m)\n",
      "\u001B[1;31mRuntimeError\u001B[0m: DataLoader worker (pid(s) 15976, 6368) exited unexpectedly"
     ]
    }
   ],
   "execution_count": 39
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "可以看到，当 DataLoader 采用 2 个进程时，由于每个进程都获取到了单独的数据集拷贝，因此会重复访问每一个样本。要避免这种情况，就需要在 DataLoader 中设置 worker_init_fn 来自定义每一个进程的数据集拷贝：\n",
    "> Jupyter Notebook 与多进程 (num_workers > 0) 处理时可能不兼容，尤其是在 Windows 操作系统中，因为 Windows 使用 \"spawn\" 作为默认的进程启动方法，而这在某些情况下与 Notebook 的执行环境不兼容。**在windows Dataloader使用多进程要在Python脚本文件中的main方法调用**"
   ],
   "id": "8b5121b38ebff0f3"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T14:05:19.152770100Z",
     "start_time": "2024-10-07T09:28:41.050441Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import math\n",
    "from torch.utils.data import get_worker_info\n",
    "\n",
    "def worker_init_fn(worker_id):\n",
    "    worker_info = get_worker_info()\n",
    "    dataset = worker_info.dataset  # the dataset copy in this worker process\n",
    "    overall_start = dataset.start\n",
    "    overall_end = dataset.end\n",
    "    # configure the dataset to only process the split workload\n",
    "    per_worker = int(math.ceil((overall_end - overall_start) / float(worker_info.num_workers)))\n",
    "    worker_id = worker_info.id\n",
    "    dataset.start = overall_start + worker_id * per_worker\n",
    "    dataset.end = min(dataset.start + per_worker, overall_end)\n",
    "\n",
    "# Worker 0 fetched [3, 4].  Worker 1 fetched [5, 6].\n",
    "print(list(DataLoader(ds, num_workers=2, worker_init_fn=worker_init_fn)))\n",
    "# With even more workers\n",
    "print(list(DataLoader(ds, num_workers=20, worker_init_fn=worker_init_fn)))"
   ],
   "id": "e93c8205284ebe6d",
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "DataLoader worker (pid(s) 9192, 23864) exited unexpectedly",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mEmpty\u001B[0m                                     Traceback (most recent call last)",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:1133\u001B[0m, in \u001B[0;36m_MultiProcessingDataLoaderIter._try_get_data\u001B[1;34m(self, timeout)\u001B[0m\n\u001B[0;32m   1132\u001B[0m \u001B[38;5;28;01mtry\u001B[39;00m:\n\u001B[1;32m-> 1133\u001B[0m     data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_data_queue\u001B[38;5;241m.\u001B[39mget(timeout\u001B[38;5;241m=\u001B[39mtimeout)\n\u001B[0;32m   1134\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m (\u001B[38;5;28;01mTrue\u001B[39;00m, data)\n",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\multiprocessing\\queues.py:114\u001B[0m, in \u001B[0;36mQueue.get\u001B[1;34m(self, block, timeout)\u001B[0m\n\u001B[0;32m    113\u001B[0m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_poll(timeout):\n\u001B[1;32m--> 114\u001B[0m         \u001B[38;5;28;01mraise\u001B[39;00m Empty\n\u001B[0;32m    115\u001B[0m \u001B[38;5;28;01melif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_poll():\n",
      "\u001B[1;31mEmpty\u001B[0m: ",
      "\nThe above exception was the direct cause of the following exception:\n",
      "\u001B[1;31mRuntimeError\u001B[0m                              Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[14], line 16\u001B[0m\n\u001B[0;32m     13\u001B[0m     dataset\u001B[38;5;241m.\u001B[39mend \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mmin\u001B[39m(dataset\u001B[38;5;241m.\u001B[39mstart \u001B[38;5;241m+\u001B[39m per_worker, overall_end)\n\u001B[0;32m     15\u001B[0m \u001B[38;5;66;03m# Worker 0 fetched [3, 4].  Worker 1 fetched [5, 6].\u001B[39;00m\n\u001B[1;32m---> 16\u001B[0m \u001B[38;5;28mprint\u001B[39m(\u001B[38;5;28mlist\u001B[39m(DataLoader(ds, num_workers\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m2\u001B[39m, worker_init_fn\u001B[38;5;241m=\u001B[39mworker_init_fn)))\n\u001B[0;32m     17\u001B[0m \u001B[38;5;66;03m# With even more workers\u001B[39;00m\n\u001B[0;32m     18\u001B[0m \u001B[38;5;28mprint\u001B[39m(\u001B[38;5;28mlist\u001B[39m(DataLoader(ds, num_workers\u001B[38;5;241m=\u001B[39m\u001B[38;5;241m20\u001B[39m, worker_init_fn\u001B[38;5;241m=\u001B[39mworker_init_fn)))\n",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:631\u001B[0m, in \u001B[0;36m_BaseDataLoaderIter.__next__\u001B[1;34m(self)\u001B[0m\n\u001B[0;32m    628\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_sampler_iter \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[0;32m    629\u001B[0m     \u001B[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001B[39;00m\n\u001B[0;32m    630\u001B[0m     \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_reset()  \u001B[38;5;66;03m# type: ignore[call-arg]\u001B[39;00m\n\u001B[1;32m--> 631\u001B[0m data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_next_data()\n\u001B[0;32m    632\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_num_yielded \u001B[38;5;241m+\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;241m1\u001B[39m\n\u001B[0;32m    633\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_dataset_kind \u001B[38;5;241m==\u001B[39m _DatasetKind\u001B[38;5;241m.\u001B[39mIterable \u001B[38;5;129;01mand\u001B[39;00m \\\n\u001B[0;32m    634\u001B[0m         \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_IterableDataset_len_called \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m \u001B[38;5;129;01mand\u001B[39;00m \\\n\u001B[0;32m    635\u001B[0m         \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_num_yielded \u001B[38;5;241m>\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_IterableDataset_len_called:\n",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:1329\u001B[0m, in \u001B[0;36m_MultiProcessingDataLoaderIter._next_data\u001B[1;34m(self)\u001B[0m\n\u001B[0;32m   1326\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_process_data(data)\n\u001B[0;32m   1328\u001B[0m \u001B[38;5;28;01massert\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_shutdown \u001B[38;5;129;01mand\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_tasks_outstanding \u001B[38;5;241m>\u001B[39m \u001B[38;5;241m0\u001B[39m\n\u001B[1;32m-> 1329\u001B[0m idx, data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_get_data()\n\u001B[0;32m   1330\u001B[0m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_tasks_outstanding \u001B[38;5;241m-\u001B[39m\u001B[38;5;241m=\u001B[39m \u001B[38;5;241m1\u001B[39m\n\u001B[0;32m   1331\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_dataset_kind \u001B[38;5;241m==\u001B[39m _DatasetKind\u001B[38;5;241m.\u001B[39mIterable:\n\u001B[0;32m   1332\u001B[0m     \u001B[38;5;66;03m# Check for _IterableDatasetStopIteration\u001B[39;00m\n",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:1295\u001B[0m, in \u001B[0;36m_MultiProcessingDataLoaderIter._get_data\u001B[1;34m(self)\u001B[0m\n\u001B[0;32m   1291\u001B[0m     \u001B[38;5;66;03m# In this case, `self._data_queue` is a `queue.Queue`,. But we don't\u001B[39;00m\n\u001B[0;32m   1292\u001B[0m     \u001B[38;5;66;03m# need to call `.task_done()` because we don't use `.join()`.\u001B[39;00m\n\u001B[0;32m   1293\u001B[0m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[0;32m   1294\u001B[0m     \u001B[38;5;28;01mwhile\u001B[39;00m \u001B[38;5;28;01mTrue\u001B[39;00m:\n\u001B[1;32m-> 1295\u001B[0m         success, data \u001B[38;5;241m=\u001B[39m \u001B[38;5;28mself\u001B[39m\u001B[38;5;241m.\u001B[39m_try_get_data()\n\u001B[0;32m   1296\u001B[0m         \u001B[38;5;28;01mif\u001B[39;00m success:\n\u001B[0;32m   1297\u001B[0m             \u001B[38;5;28;01mreturn\u001B[39;00m data\n",
      "File \u001B[1;32mD:\\Software\\anaconda3\\envs\\lzz\\Lib\\site-packages\\torch\\utils\\data\\dataloader.py:1146\u001B[0m, in \u001B[0;36m_MultiProcessingDataLoaderIter._try_get_data\u001B[1;34m(self, timeout)\u001B[0m\n\u001B[0;32m   1144\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mlen\u001B[39m(failed_workers) \u001B[38;5;241m>\u001B[39m \u001B[38;5;241m0\u001B[39m:\n\u001B[0;32m   1145\u001B[0m     pids_str \u001B[38;5;241m=\u001B[39m \u001B[38;5;124m'\u001B[39m\u001B[38;5;124m, \u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;241m.\u001B[39mjoin(\u001B[38;5;28mstr\u001B[39m(w\u001B[38;5;241m.\u001B[39mpid) \u001B[38;5;28;01mfor\u001B[39;00m w \u001B[38;5;129;01min\u001B[39;00m failed_workers)\n\u001B[1;32m-> 1146\u001B[0m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mRuntimeError\u001B[39;00m(\u001B[38;5;124mf\u001B[39m\u001B[38;5;124m'\u001B[39m\u001B[38;5;124mDataLoader worker (pid(s) \u001B[39m\u001B[38;5;132;01m{\u001B[39;00mpids_str\u001B[38;5;132;01m}\u001B[39;00m\u001B[38;5;124m) exited unexpectedly\u001B[39m\u001B[38;5;124m'\u001B[39m) \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01me\u001B[39;00m\n\u001B[0;32m   1147\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(e, queue\u001B[38;5;241m.\u001B[39mEmpty):\n\u001B[0;32m   1148\u001B[0m     \u001B[38;5;28;01mreturn\u001B[39;00m (\u001B[38;5;28;01mFalse\u001B[39;00m, \u001B[38;5;28;01mNone\u001B[39;00m)\n",
      "\u001B[1;31mRuntimeError\u001B[0m: DataLoader worker (pid(s) 9192, 23864) exited unexpectedly"
     ]
    }
   ],
   "execution_count": 14
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "下面我们以加载一个图像分类数据集为例，看看如何创建一个自定义的映射型数据集：\n",
    "```python\n",
    "import os\n",
    "import pandas as pd\n",
    "from torchvision.io import read_image\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "class CustomImageDataset(Dataset):\n",
    "    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):\n",
    "        self.img_labels = pd.read_csv(annotations_file)\n",
    "        self.img_dir = img_dir\n",
    "        self.transform = transform\n",
    "        self.target_transform = target_transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.img_labels)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])\n",
    "        image = read_image(img_path)\n",
    "        label = self.img_labels.iloc[idx, 1]\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "        if self.target_transform:\n",
    "            label = self.target_transform(label)\n",
    "        return image, label\n",
    "```\n",
    "可以看到，我们实现了 __init__()、__len__() 和 __getitem__() 三个函数，其中：\n",
    "\n",
    "- __init__() 初始化数据集参数，这里设置了图像的存储目录、标签（通过读取标签 csv 文件）以及样本和标签的数据转换函数；\n",
    "- __len__() 返回数据集中样本的个数；\n",
    "- __getitem__() 映射型数据集的核心，根据给定的索引 idx 返回样本。这里会根据索引从目录下通过 read_image 读取图片和从 csv 文件中读取图片标签，并且返回处理后的图像和标签。\n",
    "\n",
    "### DataLoaders\n",
    "前面的数据集 Dataset 类提供了一种按照索引访问样本的方式。不过在实际训练模型时，我们都需要先将数据集切分为很多的 mini-batches，然后按批 (batch) 将样本送入模型，并且循环这一过程，每一个完整遍历所有样本的循环称为一个 epoch。\n",
    "\n",
    "> 训练模型时，我们通常会在每次 epoch 循环开始前随机打乱样本顺序以缓解过拟合。\n",
    "\n",
    "Pytorch 提供了 DataLoader 类专门负责处理这些操作，除了基本的 dataset（数据集）和 batch_size （batch 大小）参数以外，还有以下常用参数：\n",
    "\n",
    "- shuffle：是否打乱数据集；\n",
    "- sampler：采样器，也就是一个索引上的迭代器；\n",
    "- collate_fn：批处理函数，用于对采样出的一个 batch 中的样本进行处理（例如前面提过的 Padding 操作）。\n",
    "例如，我们按照 batch = 64 遍历 Pytorch 自带的图像分类 FashionMNIST 数据集（每个样本是一张 28 X 28 的灰度图，以及分类标签），并且打乱数据集："
   ],
   "id": "82f94cadb8ada979"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T14:05:19.152770100Z",
     "start_time": "2024-10-07T09:48:07.442994Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor\n",
    "\n",
    "training_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")\n",
    "\n",
    "test_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")\n",
    "\n",
    "train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)\n",
    "test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)\n",
    "\n",
    "train_features, train_labels = next(iter(train_dataloader))\n",
    "print(f\"Feature batch shape: {train_features.size()}\")\n",
    "print(f\"Labels batch shape: {train_labels.size()}\")\n",
    "\n",
    "img = train_features[0].squeeze()\n",
    "label = train_labels[0]\n",
    "print(img.shape)\n",
    "print(f\"Label: {label}\")"
   ],
   "id": "770e3090ba8ea536",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz\n",
      "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-images-idx3-ubyte.gz to data\\FashionMNIST\\raw\\train-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 26421880/26421880 [00:20<00:00, 1313986.80it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting data\\FashionMNIST\\raw\\train-images-idx3-ubyte.gz to data\\FashionMNIST\\raw\n",
      "\n",
      "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz\n",
      "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/train-labels-idx1-ubyte.gz to data\\FashionMNIST\\raw\\train-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 29515/29515 [00:00<00:00, 108078.91it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting data\\FashionMNIST\\raw\\train-labels-idx1-ubyte.gz to data\\FashionMNIST\\raw\n",
      "\n",
      "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz\n",
      "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-images-idx3-ubyte.gz to data\\FashionMNIST\\raw\\t10k-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 4422102/4422102 [00:09<00:00, 488408.03it/s] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting data\\FashionMNIST\\raw\\t10k-images-idx3-ubyte.gz to data\\FashionMNIST\\raw\n",
      "\n",
      "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz\n",
      "Downloading http://fashion-mnist.s3-website.eu-central-1.amazonaws.com/t10k-labels-idx1-ubyte.gz to data\\FashionMNIST\\raw\\t10k-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5148/5148 [00:00<?, ?it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting data\\FashionMNIST\\raw\\t10k-labels-idx1-ubyte.gz to data\\FashionMNIST\\raw\n",
      "\n",
      "Feature batch shape: torch.Size([64, 1, 28, 28])\n",
      "Labels batch shape: torch.Size([64])\n",
      "torch.Size([28, 28])\n",
      "Label: 0\n"
     ]
    }
   ],
   "execution_count": 15
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "**数据加载顺序和 Sampler 类**\n",
    "\n",
    "对于迭代型数据集来说，数据的加载顺序直接由用户控制，用户可以精确地控制每一个 batch 中返回的样本，因此不需要使用 Sampler 类。\n",
    "\n",
    "对于映射型数据集来说，由于索引可以不是整数，因此我们可以通过 Sampler 对象来设置加载时的索引序列，即设置一个索引上的迭代器。如果设置了 shuffle 参数，DataLoader 就会自动创建一个顺序或乱序的 sampler，我们也可以通过 sampler 参数传入一个自定义的 Sampler 对象。\n",
    "\n",
    "常见的 Sampler 对象包括序列采样器 SequentialSampler 和随机采样器 RandomSampler，它们都通过传入待采样的数据集来创建："
   ],
   "id": "609f0e792419a0c4"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T14:05:19.152770100Z",
     "start_time": "2024-10-07T09:51:18.055165Z"
    }
   },
   "cell_type": "code",
   "source": [
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import SequentialSampler, RandomSampler\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor\n",
    "\n",
    "training_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")\n",
    "\n",
    "test_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")\n",
    "\n",
    "train_sampler = RandomSampler(training_data)\n",
    "test_sampler = SequentialSampler(test_data)\n",
    "\n",
    "train_dataloader = DataLoader(training_data, batch_size=64, sampler=train_sampler)\n",
    "test_dataloader = DataLoader(test_data, batch_size=64, sampler=test_sampler)\n",
    "\n",
    "train_features, train_labels = next(iter(train_dataloader))\n",
    "print(f\"Feature batch shape: {train_features.size()}\")\n",
    "print(f\"Labels batch shape: {train_labels.size()}\")\n",
    "test_features, test_labels = next(iter(test_dataloader))\n",
    "print(f\"Feature batch shape: {test_features.size()}\")\n",
    "print(f\"Labels batch shape: {test_labels.size()}\")"
   ],
   "id": "f724b89ce9745bcc",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Feature batch shape: torch.Size([64, 1, 28, 28])\n",
      "Labels batch shape: torch.Size([64])\n",
      "Feature batch shape: torch.Size([64, 1, 28, 28])\n",
      "Labels batch shape: torch.Size([64])\n"
     ]
    }
   ],
   "execution_count": 16
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "**批处理函数 collate_fn**\n",
    "\n",
    "批处理函数 collate_fn 负责对每一个采样出的 batch 中的样本进行处理。默认的 collate_fn 会进行如下操作：\n",
    "\n",
    "添加一个新维度作为 batch 维；\n",
    "自动地将 NumPy 数组和 Python 数值转换为 PyTorch 张量；\n",
    "保留原始的数据结构，例如输入是字典的话，它会输出一个包含同样键 (key) 的字典，但是将值 (value) 替换为 batched 张量（如何可以转换的话）。\n",
    "例如，如果样本是包含 3 通道的图像和一个整数型类别标签，即 (image, class_index)，那么默认的 collate_fn 会将这样的一个元组列表转换为一个包含 batched 图像张量和 batched 类别标签张量的元组。\n",
    "\n",
    "我们也可以传入手工编写的 collate_fn 函数以对数据进行自定义处理，例如前面我们介绍过的 padding 操作。\n",
    "\n",
    "## Train Model\n",
    "### Build Model\n",
    "我们还是以前面加载的 FashionMNIST 数据库为例，构建一个神经网络模型来完成图像分类。模型同样继承自 nn.Module 类，通过 __init__() 初始化模型中的层和参数，在 forward() 中定义模型的操作，例如："
   ],
   "id": "79d6474a25cbb1ec"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T14:05:19.152770100Z",
     "start_time": "2024-10-07T09:55:51.542135Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f'Using {device} device')\n",
    "\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(256, 10),\n",
    "            nn.Dropout(p=0.2)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits\n",
    "\n",
    "model = NeuralNetwork().to(device)\n",
    "print(model)"
   ],
   "id": "eb4e0eb616b27c91",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cpu device\n",
      "NeuralNetwork(\n",
      "  (flatten): Flatten(start_dim=1, end_dim=-1)\n",
      "  (linear_relu_stack): Sequential(\n",
      "    (0): Linear(in_features=784, out_features=512, bias=True)\n",
      "    (1): ReLU()\n",
      "    (2): Linear(in_features=512, out_features=256, bias=True)\n",
      "    (3): ReLU()\n",
      "    (4): Linear(in_features=256, out_features=10, bias=True)\n",
      "    (5): Dropout(p=0.2, inplace=False)\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "execution_count": 17
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "可以看到，我们构建的模型首先将二维图像通过 Flatten 层压成一维向量，然后经过两个带有 ReLU 激活函数的全连接隐藏层，最后送入到一个包含 10 个神经元的分类器以完成 10 分类任务。我们还通过在最终输出前添加 Dropout 层来缓解过拟合。\n",
    "\n",
    "最终我们构建的模型会输出一个 10 维向量（每一维对应一个类别的预测值），与先前介绍过的 pipeline 模型一样，这里输出的是 logits 值，我们需要再接一个 Softmax 层来计算最终的概率值。下面我们构建一个包含四个伪二维图像的 mini-batch 来进行预测："
   ],
   "id": "e2b436d487ffcb0f"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T14:05:19.152770100Z",
     "start_time": "2024-10-07T09:58:23.464560Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f'Using {device} device')\n",
    "\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(256, 10),\n",
    "            nn.Dropout(p=0.2)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits\n",
    "\n",
    "model = NeuralNetwork().to(device)\n",
    "\n",
    "X = torch.rand(4, 28, 28, device=device)\n",
    "logits = model(X)\n",
    "pred_probab = nn.Softmax(dim=1)(logits)\n",
    "print(pred_probab.size())\n",
    "y_pred = pred_probab.argmax(-1)\n",
    "print(f\"Predicted class: {y_pred}\")"
   ],
   "id": "e08b23f2ce8a765",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cpu device\n",
      "torch.Size([4, 10])\n",
      "Predicted class: tensor([6, 1, 1, 1])\n"
     ]
    }
   ],
   "execution_count": 18
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### Optimize Model Parameter\n",
    "在准备好数据、搭建好模型之后，我们就可以开始训练和测试（验证）模型了。正如前面所说，模型训练是一个迭代的过程，每一轮 epoch 迭代中模型都会对输入样本进行预测，然后对预测结果计算损失 (loss)，并求 loss 对每一个模型参数的偏导，最后使用优化器更新所有的模型参数。\n",
    "\n",
    "- 损失函数 (Loss function) 用于度量预测值与答案之间的差异，模型的训练过程就是最小化损失函数。Pytorch 实现了很多常见的损失函数，例如用于回归任务的均方误差 (Mean Square Error) nn.MSELoss、用于分类任务的负对数似然 (Negative Log Likelihood) nn.NLLLoss、同时结合了 nn.LogSoftmax 和 nn.NLLLoss 的交叉熵损失 (Cross Entropy) nn.CrossEntropyLoss 等。\n",
    "\n",
    "- 优化器 (Optimization) 使用特定的优化算法（例如随机梯度下降），通过在每一个训练阶段 (step) 减少（基于一个 batch 样本计算的）模型损失来调整模型参数。Pytorch 实现了很多优化器，例如 SGD、ADAM、RMSProp 等。\n",
    "\n",
    "每一轮迭代 (Epoch) 实际上包含了两个步骤：\n",
    "\n",
    "- 训练循环 (The Train Loop) 在训练集上进行迭代，尝试收敛到最佳的参数；\n",
    "- 验证/测试循环 (The Validation/Test Loop) 在测试/验证集上进行迭代以检查模型性能有没有提升。\n",
    "- 具体地，在训练循环中，优化器通过以下三个步骤进行优化：\n",
    "\n",
    "- 调用 optimizer.zero_grad() 重设模型参数的梯度。默认情况下梯度会进行累加，为了防止重复计算，在每个训练阶段开始前都需要清零梯度；\n",
    "- 通过 loss.backwards() 反向传播预测结果的损失，即计算损失对每一个参数的偏导；\n",
    "- 调用 optimizer.step() 根据梯度调整模型的参数。\n",
    "下面我们选择交叉熵作为损失函数、选择 AdamW 作为优化器，完整的训练循环和测试循环实现如下："
   ],
   "id": "ea5d79ed287a0a6a"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T14:05:19.153768500Z",
     "start_time": "2024-10-07T10:02:35.992589Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torchvision.transforms import ToTensor\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f'Using {device} device')\n",
    "\n",
    "training_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")\n",
    "\n",
    "test_data = datasets.FashionMNIST(\n",
    "    root=\"data\",\n",
    "    train=False,\n",
    "    download=True,\n",
    "    transform=ToTensor()\n",
    ")\n",
    "\n",
    "learning_rate = 1e-3\n",
    "batch_size = 64\n",
    "epochs = 3\n",
    "\n",
    "train_dataloader = DataLoader(training_data, batch_size=batch_size)\n",
    "test_dataloader = DataLoader(test_data, batch_size=batch_size)\n",
    "\n",
    "class NeuralNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(NeuralNetwork, self).__init__()\n",
    "        self.flatten = nn.Flatten()\n",
    "        self.linear_relu_stack = nn.Sequential(\n",
    "            nn.Linear(28*28, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(512, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(256, 10),\n",
    "            nn.Dropout(p=0.2)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        logits = self.linear_relu_stack(x)\n",
    "        return logits\n",
    "\n",
    "model = NeuralNetwork().to(device)\n",
    "\n",
    "def train_loop(dataloader, model, loss_fn, optimizer):\n",
    "    size = len(dataloader.dataset)\n",
    "    model.train()\n",
    "    for batch, (X, y) in enumerate(dataloader, start=1):\n",
    "        X, y = X.to(device), y.to(device)\n",
    "        # Compute prediction and loss\n",
    "        pred = model(X)\n",
    "        loss = loss_fn(pred, y)\n",
    "        # Backpropagation\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if batch % 100 == 0:\n",
    "            loss, current = loss.item(), batch * len(X)\n",
    "            print(f\"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]\")\n",
    "\n",
    "def test_loop(dataloader, model, loss_fn):\n",
    "    size = len(dataloader.dataset)\n",
    "    num_batches = len(dataloader)\n",
    "    test_loss, correct = 0, 0\n",
    "\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for X, y in dataloader:\n",
    "            X, y = X.to(device), y.to(device)\n",
    "            pred = model(X)\n",
    "            test_loss += loss_fn(pred, y).item()\n",
    "            correct += (pred.argmax(dim=-1) == y).type(torch.float).sum().item()\n",
    "\n",
    "    test_loss /= num_batches\n",
    "    correct /= size\n",
    "    print(f\"Test Error: \\n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \\n\")\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)\n",
    "\n",
    "for t in range(epochs):\n",
    "    print(f\"Epoch {t+1}\\n-------------------------------\")\n",
    "    train_loop(train_dataloader, model, loss_fn, optimizer)\n",
    "    test_loop(test_dataloader, model, loss_fn)\n",
    "print(\"Done!\")"
   ],
   "id": "79358ff41dd4c9ab",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using cpu device\n",
      "Epoch 1\n",
      "-------------------------------\n",
      "loss: 0.911238  [ 6400/60000]\n",
      "loss: 0.860158  [12800/60000]\n",
      "loss: 0.752084  [19200/60000]\n",
      "loss: 1.016698  [25600/60000]\n",
      "loss: 0.593356  [32000/60000]\n",
      "loss: 0.668662  [38400/60000]\n",
      "loss: 0.840665  [44800/60000]\n",
      "loss: 0.614395  [51200/60000]\n",
      "loss: 0.499799  [57600/60000]\n",
      "Test Error: \n",
      " Accuracy: 83.7%, Avg loss: 0.449362 \n",
      "\n",
      "Epoch 2\n",
      "-------------------------------\n",
      "loss: 0.612571  [ 6400/60000]\n",
      "loss: 0.679193  [12800/60000]\n",
      "loss: 0.499263  [19200/60000]\n",
      "loss: 0.800343  [25600/60000]\n",
      "loss: 0.332626  [32000/60000]\n",
      "loss: 0.554202  [38400/60000]\n",
      "loss: 0.727738  [44800/60000]\n",
      "loss: 0.543811  [51200/60000]\n",
      "loss: 0.398679  [57600/60000]\n",
      "Test Error: \n",
      " Accuracy: 84.7%, Avg loss: 0.413745 \n",
      "\n",
      "Epoch 3\n",
      "-------------------------------\n",
      "loss: 0.580985  [ 6400/60000]\n",
      "loss: 0.584226  [12800/60000]\n",
      "loss: 0.584401  [19200/60000]\n",
      "loss: 0.735826  [25600/60000]\n",
      "loss: 0.359801  [32000/60000]\n",
      "loss: 0.521579  [38400/60000]\n",
      "loss: 0.804655  [44800/60000]\n",
      "loss: 0.600007  [51200/60000]\n",
      "loss: 0.321204  [57600/60000]\n",
      "Test Error: \n",
      " Accuracy: 86.3%, Avg loss: 0.383885 \n",
      "\n",
      "Done!\n"
     ]
    }
   ],
   "execution_count": 19
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Save and Load Model\n",
    "### Save and Load Model Weights\n",
    "Pytorch 模型会将所有参数存储在一个状态字典 (state dictionary) 中，可以通过 Model.state_dict() 加载。Pytorch 通过 torch.save() 保存模型权重："
   ],
   "id": "c67ba367fde1257"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T14:05:19.153768500Z",
     "start_time": "2024-10-07T10:06:38.268193Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "import os\n",
    "\n",
    "os.mkdir(\"model\")\n",
    "torch.save(model.state_dict(), 'model/model_weights.pth')"
   ],
   "id": "a382beb8f6f3c4df",
   "outputs": [],
   "execution_count": 23
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "为了加载保存的权重，我们首先需要创建一个结构完全相同的模型实例，然后通过 Model.load_state_dict() 函数进行加载：",
   "id": "5cbd5f729f5d7af1"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-10-07T14:05:19.153768500Z",
     "start_time": "2024-10-07T10:07:24.690709Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model.load_state_dict(torch.load('model/model_weights.pth'))\n",
    "model.eval()"
   ],
   "id": "c933f4587f0a3274",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NeuralNetwork(\n",
       "  (flatten): Flatten(start_dim=1, end_dim=-1)\n",
       "  (linear_relu_stack): Sequential(\n",
       "    (0): Linear(in_features=784, out_features=512, bias=True)\n",
       "    (1): ReLU()\n",
       "    (2): Linear(in_features=512, out_features=256, bias=True)\n",
       "    (3): ReLU()\n",
       "    (4): Linear(in_features=256, out_features=10, bias=True)\n",
       "    (5): Dropout(p=0.2, inplace=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 24
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### Save and Load the Complete Model\n",
    "上面存储模型权重的方式虽然可以节省空间，但是加载前需要构建一个结构完全相同的模型实例来承接权重。如果我们希望在存储权重的同时，也一起保存模型结构，就需要将整个模型传给 torch.save() ：\n",
    "```python\n",
    "import torch\n",
    "import torchvision.models as models\n",
    "\n",
    "model = models.vgg16(pretrained=True)\n",
    "torch.save(model, 'model.pth')\n",
    "```\n",
    "这样就可以直接从保存的文件中加载整个模型（包括权重和结构）：\n",
    "```python\n",
    "model = torch.load('model.pth')\n",
    "```"
   ],
   "id": "c5157fb1f7daec3a"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
